from construct2 import *
from ip import IpAddressAdapter


class DnsStringAdapter(Adapter):
    def _encode(self, obj):
        parts = obj.split(".")
        parts.append("")
        return parts
    def _decode(self, obj):
        return ".".join(obj[:-1])

dns_record_class = Enum(UBInt16("class"),
    RESERVED = 0,
    INTERNET = 1,
    CHAOS = 3,
    HESIOD = 4,
    NONE = 254,
    ANY = 255,
)

dns_record_type = Enum(UBInt16("type"),
    IPv4 = 1,
    AUTHORITIVE_NAME_SERVER = 2,
    CANONICAL_NAME = 5,
    NULL = 10,
    MAIL_EXCHANGE = 15,
    TEXT = 16,
    X25 = 19,
    ISDN = 20,
    IPv6 = 28,
    UNSPECIFIED = 103,
    ALL = 255,
)

query_record = Struct("query_record",
    DnsStringAdapter(
        RepeatUntil(lambda obj, ctx: obj == "",
            PascalString("name")
        )
    ),
    dns_record_type,
    dns_record_class,
)

rdata = MetaField("rdata", lambda ctx: ctx["rdata_length"])

resource_record = Struct("resource_record",
    CString("name", terminator = "\xc0"),
    Padding(1, ignore = True),
    dns_record_type,
    dns_record_class,
    UBInt32("ttl"),
    UBInt16("rdata_length"),
    IfThenElse("data", lambda ctx: ctx["type"] == "IPv4",
        IpAddressAdapter(rdata),
        rdata
    )
)

dns = Struct("dns",
    UBInt16("id"),
    BitStruct("flags",
        Enum(Bit("type"),
            QUERY = 0,
            RESPONSE = 1,
        ),
        Enum(Nibble("opcode"),
            STANDARD_QUERY = 0,
            INVERSE_QUERY = 1,
            SERVER_STATUS_REQUEST = 2,
            NOTIFY = 4,
            UPDATE = 5,
        ),
        Flag("authoritive_answer"),
        Flag("truncation"),
        Flag("recurssion_desired"),
        Flag("recursion_available"),
        Padding(1),
        Flag("authenticated_data"),
        Flag("checking_disabled"),
        Enum(Nibble("response_code"),
            SUCCESS = 0,
            FORMAT_ERROR = 1,
            SERVER_FAILURE = 2,
            NAME_DOES_NOT_EXIST = 3,
            NOT_IMPLEMENTED = 4,
            REFUSED = 5,
            NAME_SHOULD_NOT_EXIST = 6,
            RR_SHOULD_NOT_EXIST = 7,
            RR_SHOULD_EXIST = 8,
            NOT_AUTHORITIVE = 9,
            NOT_ZONE = 10,
        ),
    ),
    UBInt16("question_count"),
    UBInt16("answer_count"),
    UBInt16("authority_count"),
    UBInt16("additional_count"),
    MetaRepeater(lambda ctx: ctx["question_count"],
        Rename("questions", query_record),
    ),
    MetaRepeater(lambda ctx: ctx["answer_count"],
        Rename("answers", resource_record),
    ),
    MetaRepeater(lambda ctx: ctx["authority_count"],
        Rename("authorities", resource_record),
    ),
#    MetaRepeater(lambda ctx: ctx["additional_count"],
#        Rename("additionals", resource_record),
#    ),
)


if __name__ == "__main__":
    cap1 = (
    "2624010000010000000000000377777706676f6f676c6503636f6d0000010001"
    ).decode("hex")
    
    cap2 = (
    "26 24 81 80 00 01"
    "00 05 00 06 00 06 03 77 77 77 06 67 6f 6f 67 6c"
    "65 03 63 6f 6d 00 00 01 00 01 c0 0c 00 05 00 01"
    "00 08 90 65 00 08 03 77 77 77 01 6c c0 10 c0 2c"
    "00 01 00 01 00 00 00 04 00 04 40 e9 b7 68 c0 2c"
    "00 01 00 01 00 00 00 04 00 04 40 e9 b7 93 c0 2c"
    "00 01 00 01 00 00 00 04 00 04 40 e9 b7 63 c0 2c"
    "00 01 00 01 00 00 00 04 00 04 40 e9 b7 67 c0 30"
    "00 02 00 01 00 00 a8 86 00 04 01 63 c0 30 c0 30"
    "00 02 00 01 00 00 a8 86 00 04 01 64 c0 30 c0 30"
    "00 02 00 01 00 00 a8 86 00 04 01 65 c0 30 c0 30"
    "00 02 00 01 00 00 a8 86 00 04 01 67 c0 30 c0 30"
    "00 02 00 01 00 00 a8 86 00 04 01 61 c0 30 c0 30"
    "00 02 00 01 00 00 a8 86 00 04 01 62 c0 30 c0 c0"
    "00 01 00 01 00 01 1d 0c 00 04 d8 ef 35 09 c0 d0"
    "00 01 00 01 00 00 ca 7c 00 04 40 e9 b3 09 c0 80"
    "00 01 00 01 00 00 c4 c5 00 04 40 e9 a1 09 c0 90"
    "00 01 00 01 00 00 43 91 00 04 40 e9 b7 09 c0 a0"
    "00 01 00 01 00 00 ca 7c 00 04 42 66 0b 09 c0 b0"
    "00 01 00 01 00 00 02 66 00 04 40 e9 a7 09"
    ).replace(" ", "").decode("hex")

    obj = dns.parse(cap1)
    print obj
    print repr(dns.build(obj))
    
    print "-" * 80
    
    obj = dns.parse(cap2)
    print obj
    print repr(dns.build(obj))
    
    


